# NotchFilter_详细注释.py - 陷波滤波器详细注释版本
# 作者：脑电波数据处理系统
# 功能：专门用于去除脑电信号中的工频干扰（50Hz/60Hz）

# ==================== 导入模块部分 ====================

import numpy as np                    # 导入NumPy库，用于数值计算和数组操作
from scipy import signal            # 导入SciPy信号处理模块，提供滤波器设计和应用功能
from typing import Union, List, Tuple, Optional  # 导入类型提示，提高代码可读性
import logging                       # 导入日志库，用于记录程序运行信息
import sys
import os

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
if project_root not in sys.path:
    sys.path.append(project_root)

from exception.base_exceptions import SignalProcessingException
from exception.exception_handler import global_exception_handler

logger = logging.getLogger(__name__)  # 创建日志记录器

# ==================== 陷波滤波器主类 ====================

class NotchFilter:
    """
    陷波滤波器工具类
    
    什么是陷波滤波器？
    陷波滤波器（Notch Filter）是一种特殊的带阻滤波器，专门用于去除特定频率的干扰信号。
    在脑电信号处理中，主要用于去除50Hz或60Hz的工频干扰。
    
    工作原理：
    1. 在目标频率（如50Hz）处创建一个"陷阱"，大幅衰减该频率的信号
    2. 对其他频率的信号影响最小，保持脑电信号的完整性
    3. 通过品质因子Q控制陷波的宽度，Q越大陷波越窄越精确
    
    应用场景：
    - 去除电源线路引入的50Hz/60Hz工频干扰
    - 去除荧光灯等电器设备产生的周期性干扰
    - 提高脑电信号的信噪比，改善后续分析的准确性
    
    设计特点：
    - 与service层完全解耦，可选择性使用
    - 支持实时滤波和批量滤波两种模式
    - 提供性能分析和参数调整功能
    """
    
    def __init__(self, sampling_rate: float = 512.0, notch_freq: float = 50.0, 
                 quality_factor: float = 30.0, filter_order: int = 4):
        """
        初始化陷波滤波器
        
        参数详解：
            sampling_rate: 采样率 (Hz)
                - 脑电设备的数据采集频率
                - 常见值：256Hz, 512Hz, 1000Hz
                - 必须大于陷波频率的2倍（奈奎斯特定理）
                
            notch_freq: 陷波频率 (Hz)
                - 需要去除的干扰频率
                - 中国/欧洲：50Hz（电网频率）
                - 美国/日本：60Hz（电网频率）
                
            quality_factor: 品质因子 Q
                - 控制陷波的宽度和选择性
                - Q = 中心频率 / 带宽
                - Q越大，陷波越窄，选择性越好
                - 典型值：20-50，过大可能导致不稳定
                
            filter_order: 滤波器阶数
                - 控制滤波器的陡峭程度
                - 阶数越高，过渡带越陡峭，但计算复杂度增加
                - 典型值：2-8阶
        
        设计理念：
        - 平衡滤波效果和计算效率
        - 最小化对有用信号的影响
        - 提供稳定可靠的实时处理能力
        """
        # 保存滤波器参数
        self.sampling_rate = sampling_rate      # 采样率
        self.notch_freq = notch_freq           # 陷波频率
        self.quality_factor = quality_factor    # 品质因子
        self.filter_order = filter_order       # 滤波器阶数
        
        # 设计滤波器系数
        # 这一步计算滤波器的数学系数，是滤波器的"大脑"
        self._design_filter()
        
        # 滤波器状态变量（用于实时滤波）
        # 实时滤波需要记住之前的状态，这些变量存储滤波器的"记忆"
        self.zi = None                         # 滤波器内部状态
        self.is_initialized = False            # 是否已初始化标志
        
        # 记录初始化信息
        logger.info(f"陷波滤波器初始化完成:")
        logger.info(f"  - 陷波频率: {notch_freq}Hz")
        logger.info(f"  - 品质因子: Q={quality_factor}")
        logger.info(f"  - 采样率: {sampling_rate}Hz")
        logger.info(f"  - 滤波器阶数: {filter_order}阶")
    
    def _design_filter(self):
        """
        设计陷波滤波器的核心方法
        
        这个方法是整个滤波器的核心，负责计算滤波器的数学系数。
        
        设计步骤：
        1. 频率归一化 - 将实际频率转换为数字滤波器可处理的归一化频率
        2. 计算陷波带宽 - 根据品质因子确定陷波的宽度
        3. 设计带阻滤波器 - 使用巴特沃斯滤波器设计方法
        4. 生成滤波器系数 - 得到可用于实际滤波的数学系数
        
        数学原理：
        - 归一化频率 = 实际频率 / (采样率/2)
        - 带宽 = 中心频率 / 品质因子
        - 使用巴特沃斯滤波器保证平坦的通带响应
        """
        try:
            # 第一步：计算奈奎斯特频率（采样率的一半）
            # 根据奈奎斯特定理，只能处理小于采样率一半的频率
            nyquist = self.sampling_rate / 2.0
            
            # 第二步：计算归一化的陷波频率
            # 数字滤波器需要归一化频率（0-1之间）
            normalized_freq = self.notch_freq / nyquist
            
            # 第三步：计算陷波带宽
            # 品质因子Q = 中心频率 / 带宽，所以带宽 = 中心频率 / Q
            bandwidth = self.notch_freq / self.quality_factor
            
            # 第四步：计算陷波的上下边界频率
            low_freq = (self.notch_freq - bandwidth/2) / nyquist    # 下边界
            high_freq = (self.notch_freq + bandwidth/2) / nyquist   # 上边界
            
            # 第五步：设计巴特沃斯带阻滤波器
            # 'bandstop'表示带阻（陷波），analog=False表示数字滤波器
            self.b, self.a = signal.butter(
                self.filter_order,              # 滤波器阶数
                [low_freq, high_freq],          # 阻带频率范围
                btype='bandstop',               # 带阻类型
                analog=False                    # 数字滤波器
            )
            
            # 记录设计成功信息
            logger.info(f"陷波滤波器设计成功:")
            logger.info(f"  - 阻带范围: {low_freq*nyquist:.1f}-{high_freq*nyquist:.1f}Hz")
            logger.info(f"  - 带宽: {bandwidth:.2f}Hz")
            logger.info(f"  - 系数b长度: {len(self.b)}, 系数a长度: {len(self.a)}")
            
        except Exception as e:
            # 如果设计失败，创建一个通过滤波器作为备选
            error_msg = f"陷波滤波器设计失败，陷波频率: {self.notch_freq}Hz"
            global_exception_handler(SignalProcessingException(error_msg), "NotchFilter._design_filter")
            logger.warning("使用通过滤波器作为备选方案")
            self.b, self.a = [1], [1]  # 通过滤波器：输出=输入
    
    def filter_single_value(self, value: Union[int, float]) -> float:
        """
        实时滤波单个数据点
        
        这个方法用于实时处理，每次只处理一个数据点。
        适用于实时脑电数据采集场景。
        
        实时滤波的挑战：
        1. 滤波器有"记忆"，需要保存之前的状态
        2. 初始状态的设置影响滤波效果
        3. 必须保证数值稳定性
        
        参数说明：
            value: 输入的单个信号值（可以是int或float）
            
        返回值：
            float: 滤波后的信号值
            
        工作流程：
        1. 检查是否首次使用，如果是则初始化滤波器状态
        2. 应用滤波器处理当前值
        3. 更新滤波器内部状态
        4. 返回滤波结果
        """
        try:
            # 第一步：检查滤波器是否已初始化
            if not self.is_initialized:
                # 首次使用时，需要初始化滤波器的内部状态
                # lfilter_zi计算稳态初始条件，避免启动瞬态
                self.zi = signal.lfilter_zi(self.b, self.a) * value
                self.is_initialized = True
                logger.debug(f"滤波器状态已初始化，初始值: {value}")
            
            # 第二步：应用滤波器
            # lfilter是线性滤波器，zi是内部状态
            filtered_value, self.zi = signal.lfilter(
                self.b, self.a,        # 滤波器系数
                [value],               # 输入值（必须是数组形式）
                zi=self.zi            # 滤波器内部状态
            )
            
            # 第三步：返回滤波结果
            # filtered_value是数组，取第一个元素
            result = float(filtered_value[0])
            
            # 记录调试信息（仅在调试模式下）
            if logger.isEnabledFor(logging.DEBUG):
                logger.debug(f"实时滤波: {value:.2f} -> {result:.2f}")
            
            return result
            
        except Exception as e:
            # 如果滤波失败，记录警告并返回原值
            error_msg = f"实时滤波失败，输入值: {value}"
            global_exception_handler(SignalProcessingException(error_msg), "NotchFilter.filter_single_value")
            return float(value)
    
    def filter_signal_batch(self, signal_data: Union[List, np.ndarray]) -> np.ndarray:
        """
        批量滤波信号数据
        
        这个方法用于处理一段完整的信号数据，比实时滤波效果更好。
        适用于离线分析或批量处理场景。
        
        批量滤波的优势：
        1. 可以使用零相位滤波，消除相位延迟
        2. 处理效率更高
        3. 滤波效果更稳定
        
        参数说明：
            signal_data: 输入信号数组（List或numpy数组）
            
        返回值：
            np.ndarray: 滤波后的信号数组
            
        工作流程：
        1. 将输入转换为numpy数组
        2. 检查数据有效性
        3. 应用零相位滤波
        4. 返回滤波结果
        """
        try:
            # 第一步：数据预处理
            # 确保输入是numpy数组，并转换为浮点型
            signal_array = np.array(signal_data, dtype=float)
            
            # 第二步：数据有效性检查
            if len(signal_array) == 0:
                logger.warning("输入信号为空，返回空数组")
                return signal_array
            
            if len(signal_array) < 2 * self.filter_order:
                logger.warning(f"信号长度({len(signal_array)})小于推荐长度({2*self.filter_order})")
            
            # 第三步：应用零相位滤波
            # filtfilt进行前向和后向滤波，消除相位延迟
            # 这是批量处理的核心优势
            filtered_signal = signal.filtfilt(self.b, self.a, signal_array)
            
            # 第四步：记录处理信息
            logger.info(f"批量滤波完成: 处理了{len(signal_array)}个数据点")
            
            # 计算滤波效果统计
            original_power = np.var(signal_array)
            filtered_power = np.var(filtered_signal)
            noise_reduction_db = 10 * np.log10(original_power / filtered_power) if filtered_power > 0 else 0
            
            logger.info(f"滤波效果: 功率降低 {noise_reduction_db:.2f} dB")
            
            return filtered_signal
            
        except Exception as e:
            # 如果批量滤波失败，记录错误并返回原数据
            error_msg = f"批量滤波失败，信号长度: {len(signal_data) if hasattr(signal_data, '__len__') else 'unknown'}"
            global_exception_handler(SignalProcessingException(error_msg), "NotchFilter.filter_signal_batch")
            return np.array(signal_data, dtype=float)
    
    def filter_multi_channel(self, multi_channel_data: dict) -> dict:
        """
        多通道信号滤波
        
        脑电设备通常有多个电极通道，这个方法可以同时处理多个通道的数据。
        
        应用场景：
        - 多电极脑电记录
        - 不同脑区的信号同时处理
        - 保持通道间的时间同步
        
        参数说明：
            multi_channel_data: 多通道数据字典
                格式: {'channel1': [data1], 'channel2': [data2], ...}
                
        返回值：
            dict: 滤波后的多通道数据，格式与输入相同
            
        工作流程：
        1. 遍历每个通道
        2. 对每个通道独立应用滤波
        3. 记录处理结果
        4. 返回完整的多通道结果
        """
        filtered_data = {}  # 存储滤波结果
        
        # 遍历每个通道进行处理
        for channel_name, channel_data in multi_channel_data.items():
            try:
                # 对当前通道应用批量滤波
                filtered_data[channel_name] = self.filter_signal_batch(channel_data)
                
                # 记录成功信息
                logger.debug(f"通道 {channel_name} 滤波完成，数据点数: {len(channel_data)}")
                
            except Exception as e:
                # 如果某个通道滤波失败，记录警告但继续处理其他通道
                error_msg = f"通道滤波失败，通道名: {channel_name}，原因: {str(e)}"
                channel_error = SignalProcessingException(error_msg, processing_type="multi_channel_filtering")
                global_exception_handler.handle_exception(channel_error)
                # 保存原始数据作为备选
                filtered_data[channel_name] = np.array(channel_data, dtype=float)
        
        logger.info(f"多通道滤波完成: 处理了{len(multi_channel_data)}个通道")
        return filtered_data
    
    def get_frequency_response(self, frequencies: Optional[np.ndarray] = None) -> Tuple[np.ndarray, np.ndarray]:
        """
        获取滤波器频率响应
        
        频率响应显示滤波器对不同频率信号的处理效果。
        这对于验证滤波器设计和调试非常重要。
        
        参数说明：
            frequencies: 要分析的频率点数组，默认自动生成
            
        返回值：
            tuple: (频率数组, 幅度响应数组)
            
        应用：
        - 验证陷波效果
        - 检查通带平坦度
        - 优化滤波器参数
        """
        if frequencies is None:
            # 自动生成频率点：从0.1Hz到奈奎斯特频率
            frequencies = np.logspace(-1, np.log10(self.sampling_rate/2), 1000)
        
        # 计算频率响应
        # freqs计算模拟滤波器的频率响应
        w, h = signal.freqs(self.b, self.a, worN=frequencies * 2 * np.pi)
        
        return frequencies, np.abs(h)
    
    def analyze_filter_performance(self, test_signal: Optional[np.ndarray] = None) -> dict:
        """
        分析滤波器性能
        
        这个方法通过实际测试来评估滤波器的性能，提供量化的评估指标。
        
        测试内容：
        1. 陷波频率处的衰减程度
        2. 其他频率的保真度
        3. 整体性能评分
        
        参数说明：
            test_signal: 测试信号，默认生成包含工频干扰的模拟脑电信号
            
        返回值：
            dict: 详细的性能分析结果
        """
        if test_signal is None:
            # 生成标准测试信号
            logger.info("生成标准测试信号...")
            
            # 创建10秒的时间轴
            t = np.linspace(0, 10, int(10 * self.sampling_rate))
            
            # 模拟真实的脑电信号成分
            brain_signal = (
                np.sin(2 * np.pi * 8 * t) +      # Alpha波 (8Hz)
                0.5 * np.sin(2 * np.pi * 4 * t) + # Theta波 (4Hz)
                0.3 * np.sin(2 * np.pi * 15 * t) + # Beta波 (15Hz)
                0.1 * np.random.randn(len(t))     # 随机噪声
            )
            
            # 添加工频干扰（这是我们要去除的）
            power_line_noise = 2 * np.sin(2 * np.pi * self.notch_freq * t)
            
            # 合成最终测试信号
            test_signal = brain_signal + power_line_noise
            
            logger.info(f"测试信号生成完成: {len(test_signal)}个采样点")
        
        # 应用滤波器
        filtered_signal = self.filter_signal_batch(test_signal)
        
        # 计算功率谱密度（PSD）
        freqs_orig, psd_orig = signal.periodogram(test_signal, self.sampling_rate)
        freqs_filt, psd_filt = signal.periodogram(filtered_signal, self.sampling_rate)
        
        # 分析陷波效果
        # 找到最接近陷波频率的频率点
        notch_idx = np.argmin(np.abs(freqs_orig - self.notch_freq))
        
        # 计算陷波频率处的衰减（单位：dB）
        if psd_orig[notch_idx] > 0 and psd_filt[notch_idx] > 0:
            attenuation_db = 10 * np.log10(psd_filt[notch_idx] / psd_orig[notch_idx])
        else:
            attenuation_db = -100  # 极大衰减
        
        # 分析信号保真度
        # 选择远离陷波频率的频率点
        preserve_mask = np.abs(freqs_orig - self.notch_freq) > self.notch_freq / self.quality_factor
        
        if np.any(preserve_mask) and np.any(psd_orig[preserve_mask] > 0):
            preservation_ratio = np.mean(psd_filt[preserve_mask] / psd_orig[preserve_mask])
        else:
            preservation_ratio = 1.0
        
        # 计算综合性能评分
        # 好的陷波滤波器应该：衰减大（attenuation_db很负）+ 保真度高（preservation_ratio接近1）
        performance_score = min(100, max(0, -attenuation_db * 5 + preservation_ratio * 50))
        
        # 返回详细分析结果
        result = {
            'notch_attenuation_db': attenuation_db,           # 陷波衰减（dB）
            'signal_preservation_ratio': preservation_ratio,   # 信号保真度
            'effective_bandwidth_hz': self.notch_freq / self.quality_factor,  # 有效带宽
            'filter_delay_samples': len(self.a) - 1,          # 滤波器延迟
            'performance_score': performance_score             # 综合评分（0-100）
        }
        
        # 记录分析结果
        logger.info("滤波器性能分析完成:")
        logger.info(f"  - 陷波衰减: {attenuation_db:.1f} dB")
        logger.info(f"  - 信号保真度: {preservation_ratio:.3f}")
        logger.info(f"  - 综合评分: {performance_score:.1f}/100")
        
        return result
    
    def reset_filter_state(self):
        """
        重置滤波器状态
        
        在以下情况需要重置：
        1. 开始处理新的信号序列
        2. 滤波器参数发生变化
        3. 检测到异常状态
        
        重置后，下次调用filter_single_value时会重新初始化。
        """
        self.zi = None                    # 清除内部状态
        self.is_initialized = False       # 标记为未初始化
        logger.info("陷波滤波器状态已重置")
    
    def update_notch_frequency(self, new_freq: float):
        """
        更新陷波频率并重新设计滤波器
        
        这个方法允许在运行时动态调整陷波频率。
        适用于需要适应不同工频环境的场景。
        
        参数说明：
            new_freq: 新的陷波频率（Hz）
            
        使用场景：
        - 从50Hz环境切换到60Hz环境
        - 检测到不同频率的干扰
        - 用户手动调整
        """
        logger.info(f"更新陷波频率: {self.notch_freq}Hz -> {new_freq}Hz")
        
        # 更新频率参数
        self.notch_freq = new_freq
        
        # 重新设计滤波器
        self._design_filter()
        
        # 重置滤波器状态
        self.reset_filter_state()
        
        logger.info(f"陷波频率更新完成: {new_freq}Hz")
    
    def get_filter_info(self) -> dict:
        """
        获取滤波器详细信息
        
        返回滤波器的所有重要参数和状态信息。
        用于调试、监控和用户界面显示。
        
        返回值：
            dict: 包含所有滤波器信息的字典
        """
        return {
            'sampling_rate': self.sampling_rate,              # 采样率
            'notch_frequency': self.notch_freq,               # 陷波频率
            'quality_factor': self.quality_factor,            # 品质因子
            'filter_order': self.filter_order,                # 滤波器阶数
            'bandwidth': self.notch_freq / self.quality_factor, # 陷波带宽
            'is_initialized': self.is_initialized,            # 初始化状态
            'filter_type': 'Butterworth Bandstop',           # 滤波器类型
            'coefficients_b_length': len(self.b),            # B系数长度
            'coefficients_a_length': len(self.a)             # A系数长度
        }


# ==================== 便捷函数部分 ====================

def create_notch_filter_50hz(sampling_rate: float = 512.0) -> NotchFilter:
    """
    创建50Hz陷波滤波器（适用于中国/欧洲）
    
    这是一个便捷函数，使用优化的参数创建50Hz陷波滤波器。
    
    参数说明：
        sampling_rate: 采样率，默认512Hz
        
    返回值：
        NotchFilter: 配置好的50Hz陷波滤波器对象
        
    使用示例：
        filter_50hz = create_notch_filter_50hz()
        filtered_data = filter_50hz.filter_signal_batch(raw_data)
    """
    logger.info("创建50Hz陷波滤波器（中国/欧洲标准）")
    return NotchFilter(
        sampling_rate=sampling_rate, 
        notch_freq=50.0,           # 50Hz工频
        quality_factor=30.0,       # 优化的Q值
        filter_order=4             # 4阶滤波器
    )

def create_notch_filter_60hz(sampling_rate: float = 512.0) -> NotchFilter:
    """
    创建60Hz陷波滤波器（适用于美国）
    
    这是一个便捷函数，使用优化的参数创建60Hz陷波滤波器。
    
    参数说明：
        sampling_rate: 采样率，默认512Hz
        
    返回值：
        NotchFilter: 配置好的60Hz陷波滤波器对象
        
    使用示例：
        filter_60hz = create_notch_filter_60hz()
        filtered_data = filter_60hz.filter_signal_batch(raw_data)
    """
    logger.info("创建60Hz陷波滤波器（美国标准）")
    return NotchFilter(
        sampling_rate=sampling_rate, 
        notch_freq=60.0,           # 60Hz工频
        quality_factor=30.0,       # 优化的Q值
        filter_order=4             # 4阶滤波器
    )

def quick_filter_50hz(signal_data: Union[List, np.ndarray], sampling_rate: float = 512.0) -> np.ndarray:
    """
    快速50Hz陷波滤波
    
    这是一个一步到位的便捷函数，适用于简单的批量滤波需求。
    
    参数说明：
        signal_data: 输入信号数据
        sampling_rate: 采样率
        
    返回值：
        np.ndarray: 滤波后的信号
        
    使用示例：
        filtered_data = quick_filter_50hz(raw_eeg_data)
    """
    logger.info("执行快速50Hz陷波滤波")
    filter_obj = create_notch_filter_50hz(sampling_rate)
    return filter_obj.filter_signal_batch(signal_data)

def quick_filter_60hz(signal_data: Union[List, np.ndarray], sampling_rate: float = 512.0) -> np.ndarray:
    """
    快速60Hz陷波滤波
    
    这是一个一步到位的便捷函数，适用于简单的批量滤波需求。
    
    参数说明：
        signal_data: 输入信号数据
        sampling_rate: 采样率
        
    返回值：
        np.ndarray: 滤波后的信号
        
    使用示例：
        filtered_data = quick_filter_60hz(raw_eeg_data)
    """
    logger.info("执行快速60Hz陷波滤波")
    filter_obj = create_notch_filter_60hz(sampling_rate)
    return filter_obj.filter_signal_batch(signal_data)

# ==================== 使用示例和说明 ====================

"""
陷波滤波器使用指南：

1. 基本使用：
   # 创建滤波器
   notch_filter = NotchFilter(sampling_rate=512, notch_freq=50)
   
   # 实时滤波
   filtered_value = notch_filter.filter_single_value(raw_value)
   
   # 批量滤波
   filtered_signal = notch_filter.filter_signal_batch(raw_signal)

2. 便捷使用：
   # 快速50Hz滤波
   filtered_data = quick_filter_50hz(eeg_data)
   
   # 快速60Hz滤波
   filtered_data = quick_filter_60hz(eeg_data)

3. 性能分析：
   performance = notch_filter.analyze_filter_performance()
   print(f"陷波衰减: {performance['notch_attenuation_db']:.1f} dB")

4. 参数调整：
   # 更窄的陷波（更高的Q值）
   narrow_filter = NotchFilter(quality_factor=50)
   
   # 更宽的陷波（更低的Q值）
   wide_filter = NotchFilter(quality_factor=20)

5. 多通道处理：
   multi_data = {'C3': eeg_c3, 'C4': eeg_c4, 'Cz': eeg_cz}
   filtered_multi = notch_filter.filter_multi_channel(multi_data)

注意事项：
- 陷波滤波器会引入轻微的相位延迟
- 批量滤波（filtfilt）可以消除相位延迟，但需要完整的信号
- 实时滤波适用于在线处理，但有相位延迟
- 选择合适的Q值很重要：太高可能不稳定，太低可能影响有用信号
- 确保采样率至少是陷波频率的4倍以上
"""